import os
import shutil
from tqdm import tqdm
from preprocess.file_converter import convert_file
from preprocess.file_paser import parse_file
from preprocess.file_spliter import split_file, traverse_dir
from pasers.img_parser.img_ocr import image_ocr
from rapidocr import RapidOCR
import uuid
import argparse

# 创建 ArgumentParser 对象
parser = argparse.ArgumentParser(description="命令行参数示例")
# 添加命令行参数
parser.add_argument('--input', type=str, default=r'C:\Users\Administrator\Desktop\doc_paser-master\doc_paser-master\第三批附件', help='输入文件路径')
# parser.add_argument('--input', type=str, required=True, help='输入文件路径')
parser.add_argument('--output', type=str, default='', help='输出文件路径，不传则输出文件夹与输入文件夹在同一父目录')
parser.add_argument('--chunk_size', type=int, default=800, help='切块大小')
parser.add_argument('--chunk_overlap', type=int, default=150, help='切块重叠大小')
parser.add_argument('--use_cpu_only', action="store_true", help='RapidOCR仅使用CPU')
parser.add_argument('--keep_intermediate', action="store_true", help='是否保留中间产物文件（默认删除）')

# 解析参数
args = parser.parse_args()
use_gpu = not args.use_cpu_only
ocr = RapidOCR(params={'EngineConfig.use_cuda': use_gpu, 'Global.with_onnx': True})
if not args.output:
    args.output = os.path.dirname(args.input.rstrip('/'))
SUPPORTED_FORMATS = {'doc', 'docx', 'xls', 'xlsx', 'pdf', 'ofd', 'xml','txt',
                     'jpg', 'jpeg', 'png', 'jfif', 'exif', 'gif', 'tiff', 'bmp', 'wmf'}
split_size = args.chunk_size
register_dict = {}
failed_ls = []
base_name = os.path.basename(args.input.rstrip('/'))
uuid_dir = os.path.join(args.output, f'{base_name}_uuid_files')
os.makedirs(uuid_dir, exist_ok=True)
namespace = uuid.NAMESPACE_DNS
bar = tqdm(traverse_dir(args.input, None), desc='为文件准备uuid')
for fp, fn in bar:
    file_suffix = fp.split('.')[-1]
    if file_suffix in SUPPORTED_FORMATS:
        dir_name = os.path.dirname(fp)
        uuid_name = str(uuid.uuid5(namespace, fp))
        register_dict[uuid_name] = {'path': fp, 'name': fn + '.' + file_suffix}
        shutil.copy(fp, os.path.join(uuid_dir, uuid_name + '.' + file_suffix))
    else:
        failed_ls.append({'filePath': fp, 'reason': '文件格式不支持'})

converted_dir = os.path.join(args.output, f'{base_name}_converted')
ofd_unzip_dir = os.path.join(args.output, f'{base_name}_ofd_unzipped')
image_saved_dir = os.path.join(args.output, f'{base_name}_images')
parse_out_dir = os.path.join(args.output, f'{base_name}_out')
seg_out_dir = os.path.join(args.output, f'{base_name}_seg')

convert_file(failed_ls, uuid_dir, converted_dir, converted_to_pdf=False)
parse_file(failed_ls, converted_dir, parse_out_dir, ofd_unzip_dir, image_saved_dir)
image_ocr(parse_out_dir, image_saved_dir, ocr)
split_file(failed_ls, parse_out_dir, converted_dir, seg_out_dir, split_size, args.chunk_overlap, register_dict, ocr,
           with_position=False)
print(register_dict)
# if not args.keep_intermediate:
#     shutil.rmtree(uuid_dir)
#     shutil.rmtree(converted_dir)
#     shutil.rmtree(ofd_unzip_dir)
#     shutil.rmtree(parse_out_dir)
